'''
Layer to unify h5py access for mapping.
'''
import h5py
import numpy as np

# defaults.


class H5_Driver(object):

	'''
	Opens and h5py file that overwrites anything there.
	'''
	def open_write(self, ham_file, name, dt, chunk_size, init_size):

		# Record information.
		self.ham_file = ham_file
		self.dt = dt
		self.name = name
		self.chunk_size = chunk_size
		self.init_size = init_size
			
		# Open h5 file.
		self.h5_out = h5py.File(ham_file, 'w')
		
		# Create h5py object.
		self.data = self.h5_out.create_dataset(\
			name, \
			(init_size,), \
			dt, \
			maxshape=None, \
			compression='gzip', \
			chunks=(chunk_size,) \
		)

		# Setup attributes.
		self.active_idx = 0
		self.max_idx = 0
		
	'''
	Save array to h5 object.
	This is only done on the first write.
	'''
	def save_array(self, a):
		# compute the size needed to save array.
		new_ptr = self.active_idx + a.size

		# check if we need to increate object size.
		if new_ptr > self.data.len():
			# grow h5 dataset.
			self.data.resize((new_ptr,))
			
		# save data to disk.
		self.data[self.active_idx:new_ptr] = a

		# update final pointer.
		self.active_idx = new_ptr
		self.max_idx = new_ptr

	'''
	Saves and updated piece of the array back to disk.
	'''
	def update_array(self, a):
		# note slice where this chunk was taken.
		i = self.removed_i
		j = self.removed_j
		
		# pull this chunk from object.
		self.data[i:j] = a


	'''
	Reads a specific piece of data from h5 object.
	'''
	def retrieve_array(self, i, j):
		# check bounds.
		if j > self.data.len():
			j = self.data.len()
		
		# note slice where this chunk is taken.
		self.removed_i = i
		self.removed_j = j
		
		# pull this chunk from object.
		return self.data[i:j]

	'''
	Must be called when finished with h5 object.
	Shrinks data set to active size if too large.
	'''
	def close(self):
		# check if we need to shrink.
		if self.active_idx < self.data.len():
			# shrink h5 dataset.
			self.data.resize((self.active_idx,))

		# save meta data.
		self.data.attrs['max_idx'] = self.max_idx
		
		# close the file pointer.
		self.h5_out.close()
	
